package org.gim.thirdService.AIServiceAPI.config;

import cn.hutool.core.util.ObjectUtil;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import dev.langchain4j.model.chat.StreamingChatModel;
import dev.langchain4j.model.openai.OpenAiChatModelName;
import dev.langchain4j.model.openai.OpenAiStreamingChatModel;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.service.TokenStream;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.gim.thirdService.AIServiceAPI.service.AisServices;
import org.gim.thirdService.AIServiceAPI.tools.WeatherTool;
import org.gim.thirdService.KnowledgeService.service.DocumentIngestService;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;


import javax.annotation.Resource;

import java.io.File;
import java.io.FileOutputStream;
import java.nio.file.Files;
import java.time.Duration;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

import static dev.langchain4j.model.openai.OpenAiChatModelName.*;
import static org.gim.utils.ocr.BaiduOcrUtil.recognizeMixedLang;

@Slf4j
@Component
public class ChatWebSocketHandler extends TextWebSocketHandler {

    private final UiUiConfig uiUiConfig;

    @Resource
    private WeatherTool weatherTool;

    @Value("${ufop.local-storage-path}")
    private String fileBasePath;

    @Resource
    private DocumentIngestService documentIngestService;

    @Autowired
    public ChatWebSocketHandler(UiUiConfig uiUiConfig) {
        this.uiUiConfig = uiUiConfig;
    }

    /**
     * 每个 session 对应一个 single-thread executor，保证顺序发送（避免 concurrent send 导致的 TEXT_PARTIAL_WRITING）
     */
    private final Map<String, ExecutorService> sessionExecutors = new ConcurrentHashMap<>();

    private ExecutorService getOrCreateExecutor(WebSocketSession session) {
        return sessionExecutors.computeIfAbsent(session.getId(), id ->
                Executors.newSingleThreadExecutor(r -> {
                    Thread t = new Thread(r);
                    t.setName("ws-sender-" + id);
                    t.setDaemon(true);
                    return t;
                }));
    }

    /**
     * 模型缓存：按 modelName 缓存 StreamingChatModel 实例，避免每次请求重建。
     * 使用 ConcurrentHashMap 保证线程安全
     */
    private final Map<String, StreamingChatModel> modelCache = new ConcurrentHashMap<>();


    /**
     * 模型映射（友好名 -> 底层 model 常量），便于前端传简短名
     * 如需支持更多模型，扩展此表即可
     */
    private static final Map<String, OpenAiChatModelName> MODEL_NAME_MAP = Map.ofEntries(
            Map.entry("gpt-4.1-2025-04-14", GPT_4_1_2025_04_14),
            Map.entry("gpt-4.1", GPT_4_1),
            Map.entry("gpt-4o",GPT_4_O),
            Map.entry("gpt-4", GPT_4),
            Map.entry("gpt-3.5-turbo",GPT_3_5_TURBO),
            Map.entry("o3", O3),
            Map.entry("default", GPT_4_1)
    );

    private StreamingChatModel getOrCreateModel(String modelKey, Duration timeout, boolean strictTools) {
        // 使用 friendlyName，如果没有传入，使用 "default"
        String friendlyName = modelKey != null && !modelKey.isEmpty() ? modelKey : "default";

        // 从映射表获取底层模型
        OpenAiChatModelName resolvedModel = MODEL_NAME_MAP.getOrDefault(friendlyName, MODEL_NAME_MAP.get("default"));

        // 构建模型实例
        OpenAiStreamingChatModel model = OpenAiStreamingChatModel.builder()
                    .apiKey(uiUiConfig.getApiKey())
                    .modelName(resolvedModel)
                    .baseUrl(uiUiConfig.getBaseUrl())
                    .strictTools(strictTools)
                    .timeout(timeout != null ? timeout : Duration.ofSeconds(10))
                    .build();

            log.warn("🔥 已创建真实模型实例 → {}", resolvedModel.name());
            return model;
    }


    private void shutdownExecutor(WebSocketSession session) {
        ExecutorService ex = sessionExecutors.remove(session.getId());
        if (ex != null) {
            ex.shutdownNow();
        }
    }

    @Override
    public void handleTextMessage(WebSocketSession session, TextMessage message) {
        String payload = message.getPayload();
        log.info("[Step 1] 📩 收到原始消息: {}", payload);

        try {
            ObjectMapper mapper = new ObjectMapper();
            JsonNode json = mapper.readTree(payload);

            if (json.has("type")) {
                String type = json.get("type").asText();
                String modelName = json.get("modelName").asText();
                String knowledgeCode = json.get("knowledgeCode").asText();

                log.info("需要引入知识库数据：{}",knowledgeCode);
                log.info("[Step 2] 🔍 消息类型解析: {}", type);

                switch (type) {
                    case "file" -> {
                        log.info("[Step 3] 📎 进入附件+文字处理流程");
                        handleFileMessage(session, json,modelName,knowledgeCode);
                    }
                    case "text" -> {
                        log.info("[Step 3] 📝 进入纯文字消息处理流程");
                        handleTextChat(session, payload,modelName,knowledgeCode);
                    }
                    default -> log.warn("[Step 3] ⚠️ 未知消息类型: {}", type);
                }
                return;
            }
        } catch (Exception ignore) {
            log.warn("[Step 2] ⚠️ 消息解析 JSON 失败，按纯文本处理");
        }

    }

    /**
     * 处理文件消息
     */
    private void handleFileMessage(WebSocketSession session, JsonNode json,String modelName,String knowledgeCode) {
        long startTime = System.currentTimeMillis();
        log.info("[File Step 1] 📦 开始处理文件消息");

        try {
            String fileName = json.hasNonNull("name") ? json.get("name").asText() : null;

            List<String> filePaths = new ArrayList<>();
            if (json.hasNonNull("data") && json.get("data").isArray()) {
                for (JsonNode node : json.get("data")) {
                    filePaths.add(node.asText());
                }
            } else {
                log.warn("[File Step 2] data 字段不是数组，无法处理");
                // 也可以根据实际情况改成失败或空处理
            }

            String textMessage = json.hasNonNull("text") ? json.get("text").asText() : "";

            log.info("[File Step 3] 参数解析完成 - 文件名: {}, 文件路径数: {}, 文本长度: {}",
                    fileName, filePaths.size(), textMessage.length());

            StringBuilder extractedTextBuilder = new StringBuilder();

            log.info("filePaths=={}",filePaths);

            // 遍历所有文件路径，读取文件内容
            for (String filePath : filePaths) {

                String filePathFull = fileBasePath+File.separator+filePath;

                File file = new File(filePathFull);
                log.info("[File Step 4] 遍历文件路径: {}", filePathFull);
                if (!file.exists()) {
                    log.warn("[File Step 4] 文件不存在: {}", filePath);
                    continue;
                }

                String mimeType = Files.probeContentType(file.toPath());
                if (mimeType == null) mimeType = "application/octet-stream";

                log.info("[File Step 5] 读取文件路径: {}，MIME类型: {}", filePathFull, mimeType);

                if (mimeType.startsWith("image/")) {
                    // 图片走 OCR 识别
                    log.info("[File Step 6] 🖼 开始 OCR 识别");
                    String ocrText = recognizeMixedLang(filePathFull);
                    extractedTextBuilder.append(ocrText);
                } else if (mimeType.startsWith("text/") || isTextFile(filePathFull)) {
                    // 读取文本文件内容
                    log.info("[File Step 6] 📜 读取文本文件内容");
                    String fileContent = Files.readString(file.toPath());
                    extractedTextBuilder.append(fileContent);
                } else {
                    log.warn("[File Step 6] ⚠️ 不支持的文件类型: {}", mimeType);
                }

                log.info("extractedTextBuilder=={}",extractedTextBuilder);
            }

            log.info("文件内容：{}",extractedTextBuilder);

            // 拼接最终文本内容
            String finalPrompt = textMessage;
            if (extractedTextBuilder.length() > 0) {
                finalPrompt += "\n\n附件内容如下：\n" + extractedTextBuilder.toString();
            }

            log.info("[File Step 7] 📤 最终组合消息长度: {}", finalPrompt.length());

            if (!finalPrompt.isEmpty()) {
                log.info("[File Step 8] 🚀 进入文字处理逻辑");
                handleTextChat(session, finalPrompt,modelName,knowledgeCode);
            }

        } catch (Exception e) {
            log.error("[File Error] ❌ 文件+文字消息处理失败: {}", e.getMessage(), e);
        } finally {
            long endTime = System.currentTimeMillis();
            log.info("[File Done] ⏱ 附件+文字处理耗时: {} ms", (endTime - startTime));
        }
    }

    /** 判断是否为常见的文本文件后缀 */
    private boolean isTextFile(String fileName) {
        String lower = fileName.toLowerCase();
        return lower.endsWith(".txt") || lower.endsWith(".md") ||
                lower.endsWith(".java") || lower.endsWith(".js") ||
                lower.endsWith(".json") || lower.endsWith(".xml")|| lower.endsWith(".pdf");
    }

    /**
     * 【原有】处理 AI 文本对话
     */
    public void handleTextChat(WebSocketSession session, String message,String modelName,String knowledgeCode) {

        log.info("处理 AI 文本对话信息,消息长度：{}，选择模型：{}",message.length(),modelName);

        String userMessage = String.format("""
            你是 ChatGPT，一名**专业编程技术助手**。回答应当以清晰、简洁且准确的方式传达给用户，帮助用户快速理解和解决问题。
            
            请基于以下规则和最佳实践，以 *清晰、美观、易读* 的 Markdown 格式回答用户问题：
            
            # 💡 回答原则
            - 语气友好且专业，避免机械和冷漠
            - 如问题不明确，简洁提问澄清
            - 内容真实、安全、实用，避免编造
            - 不主动透露AI身份，除非用户询问
            - 回答要尽量简练，突出重点，避免多余信息
            
            # 🛠️ 格式与风格要求（Markdown）
            - 所有代码请用三反引号 ```` ```语言 ```` 包裹，并注明语言类型（如 `js`, `java`, `vue`, `sql` 等）
            - 使用标题（#、##、###）分层展示知识结构，便于快速定位与浏览
            - 使用有序或无序列表清晰表达多点内容，注意缩进和对齐
            - 关键术语、结论等使用 **加粗** 或 *斜体* 强调
            - 在合适的地方加入 emoji（如 ✅、⚠️、📌、🚀）增强可读性和亲切感
            - 若适用，提供示例、推荐工具、常见错误、最佳实践等
            - 尽量提供“下一步建议”或“相关扩展阅读”，鼓励用户继续提问，提高互动黏性
            
            # 📌 用户问题：
            %s
            
            请**严格遵守以上要求，回答简明扼要，避免冗长。**
            """, message);

        if(!ObjectUtil.isNull(knowledgeCode) && !StringUtils.isBlank(knowledgeCode)){
            log.info("需要关联查询知识库，增加AI模型上下文，知识库编码为：{}",knowledgeCode);

            String documents = documentIngestService.queryFromQdrantAndIntegrate(knowledgeCode,userMessage);

            log.info("相关知识文档为：{}",documents);

            userMessage = String.format(
                            "要求：\n" +
                            "1. 如果知识库中有与用户问题相关的信息，请结合知识库回答。\n" +
                            "2. 如果知识库中没有相关信息，也可以根据自身知识正常回答用户问题。\n" +
                            "3. 不添加任何多余说明或无关信息。\n\n" +
                            "用户问题：%s\n" +
                            "知识库信息：%s\n" +
                            "请直接作答：",
                    userMessage, documents
            );

        }

        StreamingChatModel model = getOrCreateModel(modelName, Duration.ofSeconds(20), true);

        AisServices assistant = AiServices.builder(AisServices.class)
                .streamingChatModel(model)
                .tools(List.of(weatherTool))
                .build();

        TokenStream tokenStream = assistant.chat11(userMessage);

//        log.info("🎯 当前会话使用模型：{}", resolvedModel.name());

        // 获取或创建 session 专属 executor（single-thread）
        ExecutorService executor = getOrCreateExecutor(session);

        // 回调不直接发送，而是提交到 executor 串行发送
        tokenStream
                .onPartialResponse(partialResponse -> {
                    // 非阻塞地提交发送任务（按顺序执行）
                    executor.submit(() -> {
                        try {
                            if (session.isOpen()) {
                                session.sendMessage(new TextMessage(partialResponse));
                            } else {
                                log.warn("Session {} 已关闭，跳过发送 partialResponse", session.getId());
                            }
                        } catch (IllegalStateException ise) {
                            // 捕获 Tomcat 状态异常并记录（理论上不会出现，因为已串行）
                            log.error("IllegalStateException while sending partialResponse to session {}: {}", session.getId(), ise.getMessage(), ise);
                        } catch (Exception e) {
                            log.error("Failed to send partialResponse to session {}: {}", session.getId(), e.getMessage(), e);
                        }
                    });
                })
                .onCompleteResponse(response -> {
                    executor.submit(() -> {
                        try {
                            if (session.isOpen()) {
                                // 发送完成信号或完整结果
                                session.sendMessage(new TextMessage("[DONE]"));
                            } else {
                                log.warn("Session {} 已关闭，跳过发送 onCompleteResponse", session.getId());
                            }
                        } catch (Exception e) {
                            log.error("Failed to send onCompleteResponse to session {}: {}", session.getId(), e.getMessage(), e);
                        }
                    });
                })
                .onError(error -> {
                    executor.submit(() -> {
                        try {
                            if (session.isOpen()) {
                                session.sendMessage(new TextMessage("[ERROR]" + error.getMessage()));
                            } else {
                                log.warn("Session {} 已关闭，跳过发送 onError", session.getId());
                            }
                        } catch (Exception e) {
                            log.error("Failed to send onError to session {}: {}", session.getId(), e.getMessage(), e);
                        }
                    });
                })
                .start();
    }

    /**
     * 当连接关闭时，清理 executor
     */
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        log.info("WebSocket session {} closed. status: {}", session.getId(), status);
        shutdownExecutor(session);
        super.afterConnectionClosed(session, status);
    }

    /**
     * 可选：连接建立时的 log
     */
    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        log.info("WebSocket session {} established", session.getId());
        // lazy create executor only when needed
        super.afterConnectionEstablished(session);
    }
}
